# 60年代风格图形渲染器
# 使用AI辅助

import multiprocessing
import multiprocessing.shared_memory
import time

import numpy as np
import scipy
import PIL.Image

def worker(i,shm1_name,shm2_name,ns,i_start,i_end):
    shm1 = multiprocessing.shared_memory.SharedMemory(name=shm1_name)
    shm2 = multiprocessing.shared_memory.SharedMemory(name=shm2_name)
    
    sxy = np.ndarray((ns,2), dtype=np.float64, buffer=shm1.buf)
    us = np.ndarray((ns,1), dtype=np.float64, buffer=shm2.buf)
    
    ncam = 1.7

    N_sphere = 2
    sphere_list = np.zeros((N_sphere,4)) #x,y,z,r
    sphere_list[0,:] = np.array([0,0,-10,2])
    sphere_list[1,:] = np.array([2,2,-8,1])

    N_light = 1
    light_list = np.zeros((N_light,3)) #x,y,z
    light_list[0,:] = np.array([10,10,0])
    
    allow_normal_map = True #增加材质的凹凸感

    def get_sphere_normal(point, collided_sphere_id):
        if point.ndim != 1:
            raise ValueError('123')
            
        normal = point - sphere_list[collided_sphere_id,0:3]
        if allow_normal_map:
            normal += 0.1*(2*np.random.rand(3)-1)
        normal /= np.linalg.norm(normal)
        return normal

    def get_brightness(point,normal,view_direction=None):
        if point.ndim != 1 or normal.ndim != 1:
            raise ValueError('123')
            
        brightness = 0.2
        
        def get_highlight(light_ray, normal, view_direction):
            if view_direction.ndim != 1:
                raise ValueError('123')
            
            h = (-view_direction/np.linalg.norm(view_direction)+light_ray/np.linalg.norm(light_ray))/2
            costheta_view = h[None,:]@normal[:,None]
            costheta_view = costheta_view/(np.linalg.norm(h)*np.linalg.norm(normal))
            costheta_view = float(costheta_view[0,0])
            return costheta_view
        
        for i in range(N_light):
            light_ray = light_list[i,:] - point
            
            is_collided = is_ray_collide_with_sphere(point+0.1*light_ray/np.linalg.norm(light_ray),light_ray,use_Zbuffer=False)[0]
            if is_collided:
                continue;
            
            costheta = light_ray[None,:]@normal[:,None]/(np.linalg.norm(light_ray)*np.linalg.norm(normal))
            costheta = float(costheta[0,0])
            
            if costheta <= 0:
                continue;
            
            brightness += costheta
            
            if view_direction is not None:
                costheta_view = get_highlight(light_ray, normal, view_direction)
                if costheta_view > 0:
                    brightness += costheta*(costheta_view**32-0.5)
            
            if brightness >= 1:
                break

        brightness = min(max(float(brightness), 0), 1)
        return brightness

    def is_ray_collide_with_this_sphere(point,direction, sphere_id):
        if point.ndim != 1 or direction.ndim != 1:
            raise ValueError('123')
        
        direction = direction / np.linalg.norm(direction)
        
        is_collided = False
        collided_coord = None
        
        A = np.sum(direction**2)
        B = (2*direction[None,:]@(point[:,None] - sphere_list[sphere_id:sphere_id+1,0:3].T))[0,0]
        C = np.sum((point[None,:] - sphere_list[sphere_id:sphere_id+1,0:3]) ** 2) - sphere_list[sphere_id,3]**2
        D = B**2 - 4*A*C
        
        is_collided = D > 0
        
        if is_collided:
            t = np.zeros((2))
            t[0] = (-B+np.sqrt(D))/(2*A)
            t[1] = (-B-np.sqrt(D))/(2*A)
            
            _temp_collided_coord = direction[None,:]*t[:,None]+point[None,:] #(2*3)
            
            if t[0] < 0 and t[1] < 0:
                is_collided = False
            elif _temp_collided_coord[0,2] > _temp_collided_coord[1,2]:
                collided_coord = _temp_collided_coord[0,:]
            else:
                collided_coord = _temp_collided_coord[1,:]

        return is_collided, collided_coord

    def is_ray_collide_with_sphere(point,direction,use_Zbuffer=True):
        if point.ndim != 1 or direction.ndim != 1:
            raise ValueError('123')
        
        is_collided = False
        collided_sphere_id = 0
        collided_coord = np.zeros((3))
        
        z_buffer = -9999
        
        for i in range(N_sphere):
            _is_collided, _collided_coord = is_ray_collide_with_this_sphere(point,direction,i)
            
            if _is_collided and (_collided_coord[2] > z_buffer or use_Zbuffer == False):
                is_collided = True
                collided_sphere_id = i
                collided_coord = _collided_coord
                z_buffer = float(_collided_coord[2])
                
            if is_collided and use_Zbuffer == False:
                break
                
        return is_collided,collided_coord,collided_sphere_id
            

    def is_ray_collide_with_ground(point, direction):
        if point.ndim != 1 or direction.ndim != 1:
            raise ValueError('123')
        ground_y = -5.0
        
        # 检查方向向量的 y 分量是否为0（射线与地面平行）
        if np.isclose(direction[1], 0.0):
            # 如果射线起点已经在平面上且方向平行，这里我们简单认为不相交
            return False, None
        
        # 计算参数 t = (ground_y - point_y) / direction_y
        t = (ground_y - point[1]) / direction[1]
        
        # 检查 t 是否大于0且在正方向上相交
        if t > 0:
            # 计算交点坐标
            intersection = point + t * direction
            return True, intersection
        else:
            return False, None

    def get_ground_color(coord,view_direction):
        if coord.ndim != 1:
            raise ValueError('123')
        
        theta = np.pi/6
        rot = np.array([[np.cos(theta),0,np.sin(theta)],[0,1,0],[-np.sin(theta),0,np.cos(theta)]])
        _coord = rot @ coord[:,None]
        
        a = np.trunc(_coord[0]) % 10
        b = np.trunc(_coord[2]) % 10
        
        res = 0.5 if (a <= 4 and b <= 4) or (a > 4 and b > 4) else 1
        
        normal = np.array([0.0,1.0,0.0])
        if allow_normal_map:
            normal += 0.05*(2*np.random.rand(3)-1)
        normal /= np.linalg.norm(normal)
        brightness = get_brightness(coord,normal,view_direction)
        return res*brightness

    for i in range(i_start,i_end): 
        xs = sxy[i,0]
        ys = sxy[i,1]

        is_collided, collided_coord =  is_ray_collide_with_ground(np.array([0,0,ncam]),np.array([xs,ys,-ncam]))
        if is_collided:
            brightness = get_ground_color(collided_coord, np.array([xs,ys,-ncam]))
            us[i,0] = brightness

        is_collided, collided_coord, collided_sphere_id = is_ray_collide_with_sphere(np.array([0,0,ncam]),np.array([xs,ys,-ncam]))
        if is_collided:
            norm = get_sphere_normal(collided_coord, collided_sphere_id)
            brightness = get_brightness(collided_coord, norm, np.array([xs,ys,-ncam]))
            us[i,0] = brightness
        
    shm1.close()
    shm2.close()
    return

if __name__ == '__main__':
    
    dx = 0.0025
    L = 1

    xs,ys = np.meshgrid(np.arange(-L,L,dx),np.arange(-L,L,dx))
    n = xs.shape[0]
    sxy = np.column_stack((xs.flatten(),ys.flatten()))
    del xs,ys
    
    ns = n**2
    
    us = np.zeros((ns,1))
    
    shm1 = multiprocessing.shared_memory.SharedMemory(create=True, size=sxy.nbytes)
    shm2 = multiprocessing.shared_memory.SharedMemory(create=True, size=us.nbytes)
    shared_sxy = np.ndarray((ns,2), dtype=np.float64, buffer=shm1.buf)
    shared_us = np.ndarray((ns,1), dtype=np.float64, buffer=shm2.buf)
    shared_sxy[:,:] = sxy[:,:]
    shared_us[:,:] = 0
    
    print(ns)
    
    num_proc = 16
    process = []
    i_start = []
    i_end = []
    
    start_time = time.time()
    
    for i in range(num_proc):
        work_range = int(ns/num_proc)
        i_start.append(work_range*i)
        i_end.append(work_range*(i+1) if i < num_proc - 1 else ns)
        
        p = multiprocessing.Process(target=worker,args=(i,shm1.name,shm2.name,ns,i_start[i],i_end[i]))
        process.append(p)
        p.start()
        
    print(i_start)
    print(i_end)
    
    for p in process:
        p.join()
    
    us = shared_us.copy()
    image = np.array(us.reshape((n,n))*255, dtype = np.uint8)
    
    shm1.close()
    shm1.unlink()
    shm2.close()
    shm2.unlink()
    
    end_time = time.time()
    print(end_time-start_time)
    
    print(us.shape)
    print(us.max())
    image = PIL.Image.fromarray(image).transpose(PIL.Image.FLIP_TOP_BOTTOM)
    image.save('output_image.png')
    image.show()
